import mplfinance as mpf
import pandas as pd
import matplotlib
import json
import os

matplotlib.use('TkAgg')  # 或者使用 'Agg', 'Qt5Agg' 等
import matplotlib.pyplot as plt

# 设置中文字体以避免显示问题
plt.rcParams['font.family'] = ['simHei']  # 设置字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义文件路径
file_path = os.path.join('data', 'stock_data(1-000001).txt')

# 读取文件内容
with open(file_path, 'r', encoding='utf-8') as file:
    content = file.read()

# 提取 JSON 数据
json_data = content[content.index('(') + 1: content.rindex(')')]
data_dict = json.loads(json_data)

# 提取 K 线数据
klines = data_dict['data']['klines']

# 解析 K 线数据
data = {'Date': [], 'Open': [], 'High': [], 'Low': [], 'Close': [], 'Volume': []}

for kline in klines:
    items = kline.split(',')
    data['Date'].append(items[0])  # 日期
    data['Open'].append(float(items[1]))  # 开盘价
    data['Close'].append(float(items[2]))  # 收盘价
    data['High'].append(float(items[3]))  # 最高价
    data['Low'].append(float(items[4]))  # 最低价
    data['Volume'].append(int(items[5]))  # 成交量

# 创建 DataFrame 并设置日期为索引
df = pd.DataFrame(data)
df['Date'] = pd.to_datetime(df['Date'])
df.set_index('Date', inplace=True)

# 计算移动平均线
df['MA5'] = df['Close'].rolling(window=5).mean()
df['MA10'] = df['Close'].rolling(window=10).mean()
df['MA30'] = df['Close'].rolling(window=30).mean()
df['MA60'] = df['Close'].rolling(window=60).mean()
df['Volume_MA10'] = df['Volume'].rolling(window=10).mean()
df['Volume_MA30'] = df['Volume'].rolling(window=30).mean()

# 自定义样式
custom_style = mpf.make_mpf_style(base_mpf_style='charles', rc={'font.family': 'simHei'})

# 定义阳线和阴线的颜色
up_color = 'red'  # 阳线颜色
down_color = 'skyblue'  # 阴线颜色

# 动态生成市场颜色
market_colors = []
for open_, close in zip(df['Open'], df['Close']):
    if close >= open_:
        market_colors.append(up_color)  # 收盘价高于或等于开盘价，阳线
    else:
        market_colors.append(down_color)  # 收盘价低于开盘价，阴线

# 绘制 K 线图和成交量，同时添加移动平均线
mpf.plot(df,
         type='candle',
         volume=True,
         style=custom_style,
         title='K线图与成交量',
         ylabel='价格',
         ylabel_lower='成交量',
         mav=(5, 10, 30, 60),  # 设置移动平均线窗口
         mavcolors=['orange', 'blue', 'green', 'purple'],  # 设置移动平均线颜色
         addplot=[
             mpf.make_addplot(df['Volume_MA10'], type='line', color='orange', panel='lower', title='Volume MA10'),
             mpf.make_addplot(df['Volume_MA30'], type='line', color='blue', panel='lower', title='Volume MA30'),
         ],
         marketcolor_overrides=market_colors,  # 使用动态生成的阳线和阴线颜色
        )

# print(market_colors)
# 显示图形
plt.show()